import os
import sys
import glob
import time
import pathlib
import webbrowser as web
from PyQt5.QtGui import  *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
import torch
from PIL import Image
from model import MattingRefine
from torch import nn
from torch.nn import functional as F
from dataset import VideoDataset, ZipDataset
from dataset import augmentation as A
from torchvision import transforms as T
from torchvision import transforms
from torchvision.transforms.functional import to_pil_image

device = torch.device('cuda')
precision = torch.float32
loader = transforms.Compose([
    transforms.ToTensor()])



class MattingProgressBar(QWidget):
    close_Signal = pyqtSignal(int)
    def __init__(self):
        super(MattingProgressBar,self).__init__()
        self.setStyleSheet(
            'QWidget{background-color:rgb(75,75,75);color:rgb(200,200,200)}')

        self.setWindowTitle('生成Alpha')
        self.setWindowFlags(Qt.WindowStaysOnTopHint)
        self.setMinimumWidth(500)
        self.mainLayout=QVBoxLayout()
        self.mainLayout.setContentsMargins(15,12,15,15)
        self.tip_layout =QHBoxLayout()
        self.tip_layout.setAlignment(Qt.AlignCenter)
        self.tip_label=QLabel("文件名：")
        self.tip_name =QLabel('xxx')
        self.tip_layout.addWidget(self.tip_label)
        self.tip_layout.addWidget(self.tip_name)
        self.frame_layout =QHBoxLayout()
        self.frame_layout.setAlignment(Qt.AlignCenter)
        self.all_frame_label = QLabel('总帧数: ')
        self.all_frame = QLabel('0')
        self.frame_layout.addWidget(self.all_frame_label)
        self.frame_layout.addWidget(self.all_frame)
        self.current_frame_layout =QHBoxLayout()
        self.current_frame_layout.setAlignment(Qt.AlignCenter)
        self.current_frame_label =QLabel('当前帧:')
        self.current_frame = QLabel('0')
        self.current_frame_layout.addWidget(self.current_frame_label)
        self.current_frame_layout.addWidget(self.current_frame)
        self.progress_layout=QHBoxLayout()
        self.p=QProgressBar()


        self.mainLayout.addLayout(self.tip_layout)
        self.mainLayout.addLayout(self.frame_layout)
        self.mainLayout.addLayout(self.current_frame_layout)
        self.mainLayout.addWidget(self.p)
        self.setLayout(self.mainLayout)

    def setValue(self,value):
        self.p.setValue(value)

    def closeEvent(self, QCloseEvent):
        self.close_Signal.emit(1)


class MattingThread(QThread):
    back_Signal =pyqtSignal(dict)
    def __init__(self,data_info):
        super(MattingThread,self).__init__()

        self.model_path =data_info['model_path']
        self.sample_type = data_info['sample_type']
        self.model_type = data_info['model_type']
        self.current_filename = data_info['filename'] #当前正在转化的文件名
        self.type = data_info['type']
        self.src_path = data_info['src_path']
        self.bg_path = data_info['bg_path']
        self.out_dir = data_info['out_dir']
        self.running =False
        self.count = 0   #当前执行帧数


        self.model = MattingRefine(backbone=self.model_type,
                              backbone_scale=0.25,
                              refine_mode=self.sample_type)
        self.model.load_state_dict(torch.load(self.model_path))
        print('start 02')
        self.model = self.model.eval().to(precision).to(device)
        print('start 03')
        self.fg_imgs,self.bg_imgs,self.all_frames = self.get_imgs()
        print('start 04')

        print('self.all_frames:',self.all_frames)


    def get_imgs(self):
        if self.type == 'video_video':
            fg= VideoDataset(self.src_path)
            bg= VideoDataset(self.bg_path)
            all_frames = fg.frame_count
        elif self.type in ['video_img','video_imgs']:
            fg= VideoDataset(self.src_path)
            bg = self.bg_path
            all_frames = fg.frame_count
        elif self.type == 'img_img':
            fg= self.src_path
            bg = self.bg_path
            all_frames = 1
        elif self.type == 'imgs_video':
            fg= self.src_path
            bg = VideoDataset(self.bg_path)
            all_frames = len(self.src_path)
        elif self.type == 'imgs_imgs':
            fg= self.src_path
            bg = self.bg_path
            all_frames = len(fg)


        return fg,bg,all_frames

    def image2tensor(self,image_name):
        image = Image.open(image_name).convert('RGB')
        image = loader(image).unsqueeze(0)
        return image.to(device, torch.float)

    def pil2tensor(self,pil_img):
        image = loader(pil_img).unsqueeze(0)
        return image.to(device, torch.float)

    def writer(self,img, path):
        img = to_pil_image(img[0].cpu())
        img.save(path)


    def video_video(self,filename,model,fg_imgs,bg_imgs):
        """
        视频前景与视频背景
        """
        with torch.no_grad():
            for i in range(fg_imgs.frame_count):
                if self.running:
                    src_tensor= self.pil2tensor(fg_imgs[i])
                    bg_tensor = self.pil2tensor(bg_imgs[i])
                    pha, fgr = model(src_tensor, bg_tensor)[:2]
                    out_dir = os.path.join(self.out_dir,filename)
                    if not os.path.exists(out_dir):
                        os.makedirs(out_dir)
                    out_path = os.path.join(out_dir,'{}.{}.png'.format(self.model_type,'%04d' % (i+1)))
                    self.writer(pha, out_path)
                    self.count += 1
                    print('{}完成'.format(i + 1))

    def video_img(self,filename,model,fg_imgs,bg_path):
        """
        单视频前景与单背景图输出
        """
        bg_tensor = self.image2tensor(bg_path)
        with torch.no_grad():
            for i in range(fg_imgs.frame_count):
                if self.running:
                    src_tensor= self.pil2tensor(fg_imgs[i])
                    pha, fgr = model(src_tensor, bg_tensor)[:2]
                    out_dir = os.path.join(self.out_dir,filename)
                    if not os.path.exists(out_dir):
                        os.makedirs(out_dir)
                    out_path = os.path.join(out_dir,'{}.{}.png'.format(self.model_type,'%04d' % (i+1)))
                    self.writer(pha, out_path)
                    self.count += 1
                    print('{}完成'.format(i + 1))

    def video_imgs(self,filename,model,fg_imgs,bg_paths):
        """
        单视频前景与单背景序列输出
        """

        with torch.no_grad():
            for i in range(fg_imgs.frame_count):
                print('self.running:',self.running)
                if self.running:
                    src_tensor= self.pil2tensor(fg_imgs[i])
                    bg_tensor =self.image2tensor(bg_paths[i])
                    pha, fgr = model(src_tensor, bg_tensor)[:2]
                    out_dir = os.path.join(self.out_dir,filename)
                    if not os.path.exists(out_dir):
                        os.makedirs(out_dir)
                    out_path = os.path.join(out_dir,'{}.{}.png'.format(self.model_type,'%04d' % (i+1)))
                    self.writer(pha, out_path)
                    self.count += 1
                    print('{}完成'.format(i + 1))

    def img_img(self,filename,model,src_path,bg_path):
        """
        单图片前景与单图片背景输出
        """
        src_tensor = self.image2tensor(src_path)
        bg_tensor = self.image2tensor(bg_path)
        with torch.no_grad():
            pha, fgr = model(src_tensor, bg_tensor)[:2]
            out_dir = os.path.join(self.out_dir, filename)
            if not os.path.exists(out_dir):
                os.makedirs(out_dir)
            out_path = os.path.join(out_dir, '{}.png'.format(self.model_type))
            self.writer(pha, out_path)
            self.count += 1
            print('完成')

    def imgs_video(self,filename,model,src_paths,bg_imgs):
        """
        序列前景与单视频背景
        """

        with torch.no_grad():
            for i in range(len(src_paths)):
                if self.running:
                    src_tensor = self.image2tensor(src_paths[i])
                    bg_tensor = self.pil2tensor(bg_imgs[i])
                    pha, fgr = model(src_tensor, bg_tensor)[:2]
                    out_dir = os.path.join(self.out_dir, filename)
                    if not os.path.exists(out_dir):
                        os.makedirs(out_dir)
                    out_path = os.path.join(out_dir, '{}.{}.png'.format(self.model_type, '%04d' % (i+1)))
                    self.writer(pha, out_path)
                    self.count += 1
                    print('{}完成'.format(i + 1))

    def imgs_imgs(self,filename,model,src_paths,bg_paths):
        """
        序列前景与序列背景
        """
        with torch.no_grad():
            for i in range(len(src_paths)):
                if self.running:
                    src_tensor = self.image2tensor(src_paths[i])
                    bg_tensor = self.image2tensor(bg_paths[i])
                    pha, fgr = model(src_tensor, bg_tensor)[:2]
                    out_dir = os.path.join(self.out_dir, filename)
                    if not os.path.exists(out_dir):
                        os.makedirs(out_dir)
                    out_path = os.path.join(out_dir, '{}.{}.png'.format(self.model_type, i + 1))
                    self.writer(pha, out_path)
                    self.count += 1
                    print('{}完成'.format(i + 1))


    def run(self):
        print('start')
        print(self.model_type)
        print(self.sample_type)

        try:
            eval('self.{}'.format(self.type))(self.current_filename,self.model,self.fg_imgs,self.bg_imgs)
            print('start 04')
        except Exception as run_ERR:
            print('run_ERR:',str(run_ERR))




class BackgroundMatting(QWidget):

    def __init__(self):
        super(BackgroundMatting, self).__init__()
        icon=QIcon()
        cgai_icon = str(pathlib.Path('CGAI.png').resolve())
        icon.addPixmap(QPixmap(cgai_icon))

        project_dir =pathlib.Path().resolve()
        model_dir = project_dir.joinpath('model')
        print('model_dir:',model_dir)
        self.mobilenetv2_model = str(model_dir.joinpath('pytorch_mobilenetv2.pth'))
        self.resnet50_model = str(model_dir.joinpath('pytorch_resnet50.pth'))
        self.resnet101_model = str(model_dir.joinpath('pytorch_resnet101.pth'))



        p=self.palette()
        p.setColor(QPalette.Base,QColor('#1C1C1C'))
        p.setColor(QPalette.Window, QColor('#393939'))
        p.setColor(QPalette.WindowText, QColor('#E8E8E8'))
        p.setColor(QPalette.Text,QColor('#1C1C1C'))


        self.setPalette(p)

        self.big_font =QFont('',20,65)
        self.mid_font =QFont('',11,75)

        self.link_btn_style='''QPushButton{color:black}
                                  QPushButton:hover{color:#FF7F24}
                                  QPushButton{background-color:#CFCFCF}
                                  QPushButton{border:2px}
                                  QPushButton{border-radius:10px}
                                  QPushButton{padding:5px 1px}'''

        self.file_btn_style='''QPushButton{color:black}
                                  QPushButton:hover{color:#FF7F24}
                                  QPushButton{background-color:#CFCFCF}
                                  QPushButton{border:2px}
                                  QPushButton{border-radius:3px}
                                  QPushButton{padding:5px 1px}'''


        self.export_btn_style='''QPushButton{color:black}
                                  QPushButton:hover{color:#FF7F24}
                                  QPushButton{background-color:#CFCFCF}
                                  QPushButton{border:2px}
                                  QPushButton{border-radius:3px}
                                  QPushButton{padding:5px 1px}'''


        self.radio_btn_style = ''' QRadioButton:hover{color:#FF7F24}
                                   QRadioButton{color : #E8E8E8}
                                   QRadioButton{border:2px}
                                   QRadioButton{border-radius:10px}
                                   QRadioButton{padding:5px 1px}'''


        self.input_type = '图片'
        self.sample_type ='sampling'
        self.model_type = 'mobilenetv2'

        self.setWindowIcon(icon)
        self.setWindowTitle('BackgroundMattingV2')
        self.setMinimumHeight(500)
        self.setMaximumWidth(650)
        self.setMaximumHeight(600)

        self.main_layout= QVBoxLayout()
        self.main_layout.setAlignment(Qt.AlignTop)
        self.main_layout.setSpacing(18)


        self.cgai_layout =QHBoxLayout()
        self.cgai_layout.setAlignment(Qt.AlignLeft)
        self.cgai_icon =QLabel()
        cgai_pixmap =QPixmap(cgai_icon)
        scaled_pixmap = cgai_pixmap.scaled(80, 80, Qt.KeepAspectRatio)  #
        self.cgai_icon.setPixmap(scaled_pixmap)  # scale_pixmap
        self.cgai_label = QLabel('BackgroundMattingV2 模型应用')
        self.cgai_label.setFont(self.big_font)
        self.cgai_layout.addWidget(self.cgai_icon)
        self.cgai_layout.addWidget(self.cgai_label)

        self.git_layout =QHBoxLayout()
        self.git_layout.setContentsMargins(0,0,0,20)
        self.gitee_btn =QPushButton('Gitee 源码')
        self.gitee_btn.clicked.connect(self._open_gitee)
        self.gitee_btn.setStyleSheet(self.link_btn_style)
        self.github_btn =QPushButton('Github 原项目')
        self.github_btn.clicked.connect(self._open_github)
        self.github_btn.setStyleSheet(self.link_btn_style)
        self.baidu_btn =QPushButton('百度云:提取码CGAI')
        self.baidu_btn.clicked.connect(self._open_baidupan)
        self.baidu_btn.setStyleSheet(self.link_btn_style)
        self.help_btn =QPushButton('使用帮助')
        self.help_btn.clicked.connect(self._help)
        self.help_btn.setStyleSheet(self.link_btn_style)
        self.git_layout.addWidget(self.gitee_btn)
        self.git_layout.addWidget(self.github_btn)
        self.git_layout.addWidget(self.baidu_btn)
        self.git_layout.addWidget(self.help_btn)

        # self.type_frame =QFrame()
        # self.type_layout= QHBoxLayout()
        # self.type_layout.setSpacing(55)
        # self.type_label = QLabel('输入类型:')
        # self.type_label.setFont(self.mid_font)
        # self.type_layout.setAlignment(Qt.AlignLeft)
        # self.video_radio= QRadioButton('视频')
        # self.video_radio.clicked.connect(self._type_changed)
        # self.img_radio= QRadioButton('图片')
        # self.img_radio.setChecked(True)
        # self.img_radio.clicked.connect(self._type_changed)
        # self.type_layout.addWidget(self.type_label)
        # self.type_layout.addWidget(self.video_radio)
        # self.type_layout.addWidget(self.img_radio)
        # self.type_frame.setLayout(self.type_layout)

        self.sample_frame=QFrame()
        self.sample_layout =QHBoxLayout()
        self.sample_layout.setSpacing(55)
        self.sample_layout.setAlignment(Qt.AlignLeft)
        self.sample_label =QLabel('采样类型:')
        self.sample_label.setFont(self.mid_font)
        self.sample_radio =QRadioButton('sampling')
        self.sample_radio.setStyleSheet(self.radio_btn_style)
        self.sample_radio.setChecked(True)
        self.sample_radio.clicked.connect(self._sample_changed)
        self.full_radio = QRadioButton('full')
        self.full_radio.setStyleSheet(self.radio_btn_style)
        self.full_radio.clicked.connect(self._sample_changed)
        self.sample_layout.addWidget(self.sample_label)
        self.sample_layout.addWidget(self.sample_radio)
        self.sample_layout.addWidget(self.full_radio)
        self.sample_frame.setLayout(self.sample_layout)

        self.model_frame= QFrame()
        self.model_layout =QHBoxLayout()
        self.model_layout.setSpacing(55)
        self.model_layout.setAlignment(Qt.AlignLeft)
        self.model_label =QLabel('选择模型:')
        self.model_label.setFont(self.mid_font)
        self.mobilenetv2_radio =QRadioButton('mobilenetv2')
        self.mobilenetv2_radio.setStyleSheet(self.radio_btn_style)
        self.mobilenetv2_radio.setChecked(True)
        self.mobilenetv2_radio.clicked.connect(self._model_changed)
        self.resnet50_radio =QRadioButton('resnet50')
        self.resnet50_radio.setStyleSheet(self.radio_btn_style)
        self.resnet50_radio.clicked.connect(self._model_changed)
        self.resnet101_radio =QRadioButton('resnet101')
        self.resnet101_radio.setStyleSheet(self.radio_btn_style)
        self.resnet101_radio.clicked.connect(self._model_changed)
        self.model_layout.addWidget(self.model_label)
        self.model_layout.addWidget(self.mobilenetv2_radio)
        self.model_layout.addWidget(self.resnet50_radio)
        self.model_layout.addWidget(self.resnet101_radio)
        self.model_frame.setLayout(self.model_layout)


        self.src_layout =QHBoxLayout()
        self.src_layout.setContentsMargins(0,25,0,0)
        self.src_label =QLabel('视频/图片目录 :')
        self.src_label.setFont(self.mid_font)
        self.src_edit =QLineEdit()
        self.src_btn= QPushButton('··')
        self.src_btn.setStyleSheet(self.file_btn_style)
        self.src_btn.clicked.connect(self._select_src)
        self.open_src_btn =QPushButton('👆')
        self.open_src_btn.setStyleSheet(self.file_btn_style)
        self.open_src_btn.clicked.connect(self._open_src)
        self.src_layout.addWidget(self.src_label)
        self.src_layout.addWidget(self.src_edit)
        self.src_layout.addWidget(self.src_btn)
        self.src_layout.addWidget(self.open_src_btn)
        self.bg_layout =QHBoxLayout()
        self.bg_label =QLabel('背景视/图目录 :')
        self.bg_label.setFont(self.mid_font)
        self.bg_edit =QLineEdit()
        self.bg_btn =QPushButton('··')
        self.open_bg_btn =QPushButton('👆')
        self.open_bg_btn.setStyleSheet(self.file_btn_style)
        self.open_bg_btn.clicked.connect(self._open_bg)
        self.bg_btn.setStyleSheet(self.file_btn_style)
        self.bg_btn.clicked.connect(self._select_bg)
        self.bg_layout.addWidget(self.bg_label)
        self.bg_layout.addWidget(self.bg_edit)
        self.bg_layout.addWidget(self.bg_btn)
        self.bg_layout.addWidget(self.open_bg_btn)
        self.export_layout =QHBoxLayout()
        self.export_label= QLabel('Alpha输出目录:')
        self.export_label.setFont(self.mid_font)
        self.export_edit =QLineEdit()
        self.export_btn =QPushButton('··')
        self.export_btn.setStyleSheet(self.file_btn_style)
        self.export_btn.clicked.connect(self._export_dir)
        self.open_export_btn =QPushButton('👆')
        self.open_export_btn.setStyleSheet(self.file_btn_style)
        self.open_export_btn.clicked.connect(self._open_export)
        self.export_layout.addWidget(self.export_label)
        self.export_layout.addWidget(self.export_edit)
        self.export_layout.addWidget(self.export_btn)
        self.export_layout.addWidget(self.open_export_btn)

        self.matte_layout =QHBoxLayout()
        self.matte_layout.setContentsMargins(0,15,0,20)
        self.matte_btn =QPushButton('生成Alpha')
        self.matte_btn.setMaximumWidth(100)
        self.matte_btn.setStyleSheet(self.export_btn_style)
        self.matte_btn.clicked.connect(self._create)
        self.matte_layout.addWidget(self.matte_btn)

        self.main_layout.addLayout(self.cgai_layout)
        self.main_layout.addLayout(self.git_layout)
        # self.main_layout.addWidget(self.type_frame)
        self.main_layout.addWidget(self.sample_frame)
        self.main_layout.addWidget(self.model_frame)
        self.main_layout.addLayout(self.src_layout)
        self.main_layout.addLayout(self.bg_layout)
        self.main_layout.addLayout(self.export_layout)
        self.main_layout.addLayout(self.matte_layout)


        self.setLayout(self.main_layout)


        self.timer = QBasicTimer()

        self.prog = MattingProgressBar()
        self.prog.close_Signal.connect(self._close_prog)


    def _type_changed(self):
        self.input_type=self.sender().text()
        # print('self.input_type:',self.input_type)

    def _sample_changed(self):
        self.sample_type=self.sender().text()
        # print('self.sample_type:',self.sample_type)

    def _model_changed(self):
        self.model_type=self.sender().text()
        # print('self.model_type:',self.model_type)

    def _select_src(self):
        dir_path = QFileDialog.getExistingDirectory(self, '选择文件夹')
        # print('dir_path:',dir_path)
        if dir_path:
            self.src_edit.setText(dir_path)

    def _open_src(self):
        os.startfile(self.src_edit.text())



    def _select_bg(self):
        dir_path = QFileDialog.getExistingDirectory(self, '选择文件夹')
        if dir_path:
            self.bg_edit.setText(dir_path)
        # bg_path = QFileDialog.getOpenFileName(self,'选择文件','','imgs(*.jpg;*.png;*.jpeg')
        # # print('bg_path:',bg_path)
        # if bg_path:
        #     self.bg_edit.setText(bg_path[0])

    def _open_bg(self):
        os.startfile(self.bg_edit.text())


    def _export_dir(self):
        dir_path = QFileDialog.getExistingDirectory(self, '选择文件夹')
        if dir_path:
            self.export_edit.setText(dir_path)

    def _open_export(self):
        os.startfile(self.export_edit.text())


    def _open_gitee(self):
        web.open(r'https://gitee.com/cgai/BackgroundMattingV2')

    def _open_github(self):
        web.open(r'https://github.com/PeterL1n/BackgroundMattingV2')

    def _open_baidupan(self):
        web.open(r'https://pan.baidu.com/s/1KJsq8p7Xr-8eRL4FVc5I1w')


    def get_model_path(self,model_type):
        model_path = ''
        if model_type == 'mobilenetv2':
            model_path = self.mobilenetv2_model
        elif model_type == 'resnet50':
            model_path = self.resnet50_model
        elif model_type == 'resnet101':
            model_path = self.resnet101_model
        return  model_path

    def get_sorted_imgs(self,dir_path):
        imgs = sorted([*glob.glob(os.path.join(dir_path, '**', '*.jpg'), recursive=True),
                            *glob.glob(os.path.join(dir_path, '**', '*.jpeg'), recursive=True),
                            *glob.glob(os.path.join(dir_path, '**', '*.tif'), recursive=True),
                            *glob.glob(os.path.join(dir_path, '**', '*.png'), recursive=True),
                            *glob.glob(os.path.join(dir_path, '**', '*.exr'), recursive=True)]
                           )
        return imgs

    def check_fg_bg(self,src_dir,bg_dir):
        """
        检查输入的前景与背景类型
        """
        result={}
        files = os.listdir(src_dir)
        bg_files = os.listdir(bg_dir)
        filename, filetype = os.path.splitext(files[0])
        bg_filename, bg_filetype = os.path.splitext(bg_files[0])
        if len(files)==1 and len(bg_files)==1:  #单视频或者单图片

            if filetype in ['.mp4', '.MP4','.mov','.MOV','.avi','.AVI','.flv','.FLV'] and bg_filetype in ['.mp4', '.MP4','.mov','.MOV','.avi','.AVI','.flv','.FLV']:
                print('单视频前景与单视频背景')#filename,model,src_path,bg_path
                src_path = os.path.join(src_dir,files[0])
                bg_path = os.path.join(bg_dir,bg_files[0])
                result['src_path'] = src_path
                result['bg_path'] = bg_path
                result['filename'] = filename
                result['type'] = 'video_video'

            elif filetype in ['.mp4', '.MP4','.mov','.MOV','.avi','.AVI','.flv','.FLV'] and bg_filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG']:
                print('单视频单背景输出')
                src_path = os.path.join(src_dir,files[0])
                bg_path = os.path.join(bg_dir,bg_files[0])
                result['src_path'] = src_path
                result['bg_path'] = bg_path
                result['filename'] = filename
                result['type'] = 'video_img'


            elif filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'] and bg_filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG']:
                print('单前景图片与单背景图片输出')
                src_path = os.path.join(src_dir,files[0])
                bg_path = os.path.join(bg_dir,bg_files[0])
                result['src_path'] = src_path
                result['bg_path'] = bg_path
                result['filename'] = filename
                result['type'] = 'img_img'

        elif len(files) == 1 and len(bg_files)>1:
            if filetype in ['.mp4', '.MP4', '.mov', '.MOV', '.avi', '.AVI', '.flv', '.FLV']  and bg_filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG']:
                print('单视频前景与序列背景输出')
                src_path = os.path.join(src_dir,files[0])
                bg_imgs =self.get_sorted_imgs(bg_dir)
                bg_path =bg_imgs
                result['src_path'] = src_path
                result['bg_path'] = bg_path
                result['filename'] = filename
                result['type'] = 'video_imgs'

        elif len(files)>1 and len(bg_files)==1:
            if filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'] and bg_filetype in ['.mp4', '.MP4','.mov','.MOV','.avi','.AVI','.flv','.FLV']:
                print('前景序列与单背景视频')
                src_imgs = self.get_sorted_imgs(src_dir)
                bg_path = os.path.join(bg_dir,bg_files[0])
                src_path = src_imgs
                result['src_path'] = src_path
                result['bg_path'] = bg_path
                result['filename'] = filename
                result['type'] = 'imgs_video'

            if filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'] and bg_filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG']:
                print('前景序列与单背景图')
                src_imgs = self.get_sorted_imgs(src_dir)
                bg_path = os.path.join(bg_dir,bg_files[0])
                src_path =src_imgs
                result['src_path'] = src_path
                result['bg_path'] = bg_path
                result['filename'] = filename
                result['type'] = 'imgs_img'

            if filetype in ['.mp4', '.MP4', '.mov', '.MOV', '.avi', '.AVI', '.flv', '.FLV']:
                print('暂不支持多视频输入!')

        elif len(files)>1 and len(bg_files)>1:
            if filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'] and bg_filetype in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG']:
                print('前景序列与背景序列')
                src_imgs = self.get_sorted_imgs(src_dir)
                bg_imgs = self.get_sorted_imgs(bg_dir)

                src_path = src_imgs
                bg_path = bg_imgs
                result['src_path'] = src_path
                result['bg_path'] = bg_path
                result['filename'] = filename
                result['type'] = 'imgs_imgs'

        else:
            print('暂不支持多视频输入')

        return result


    def _create(self):
        src_dir = self.src_edit.text()
        bg_dir = self.bg_edit.text()
        out_dir = self.export_edit.text()
        if src_dir and bg_dir and out_dir:
            model_path = self.get_model_path(self.model_type)
            #model_path,sample_type,model_type,src_dir,bg_dir,out_dir
            data_info=self.check_fg_bg(src_dir,bg_dir)
            data_info['model_path'] =model_path
            data_info['sample_type'] =self.sample_type
            data_info['model_type'] =self.model_type
            data_info['out_dir'] =out_dir
            print(data_info)
            self.matte_btn.setText('生成中...')

            self.matting_thread = MattingThread(data_info)
            self.matting_thread.running=True
            self.timer.start(200, self)
            self.matting_thread.start()
            self.prog.show()
        else:
            QMessageBox.information(None,'提升','请输入完整路径')


    def timerEvent(self, event):
        if self.matting_thread.running: #tip_name all_frame current_frame
            self.prog.tip_name.setText(self.matting_thread.current_filename)
            self.prog.all_frame.setText(str(self.matting_thread.all_frames))
            pro_value = int(self.matting_thread.count*100.0/self.matting_thread.all_frames)
            self.prog.setValue(pro_value)
            self.prog.current_frame.setText(str(self.matting_thread.count))
            if pro_value == 100:
                self.matting_thread.running=False
                self.timer.stop()
                self.prog.close()
                self.matte_btn.setText('生成Alpha')

    def _close_prog(self, close_prog):
        if close_prog:
            self.timer.stop()
            self.matting_thread.running=False
            self.matte_btn.setText('生成Alpha')

    def _help(self):
        QMessageBox.about(None,'帮助','1.采样类型只提供sampling与full,建议使用sampling就够了\n\n'
                                    '2.低分辨率(非4K HD)默认使用mobilenetv2,4K,HD推荐使用残差resnet模型\n\n'
                                    '3.支持前景与背景的格式有: \n'
                                    '   单视频前景与单视频背景(video_video)\n'
                                    '   单视频单背景输出(video_img)\n'
                                    '   单前景图片与单背景图片输出(img_img)\n'
                                    '   单视频前景与序列背景输出(video_imgs)\n'
                                    '   前景序列与单背景视频(imgs_video)\n'
                                    '   前景序列与单背景图(imgs_img)\n'
                                    '   前景序列与背景序列(imgs_imgs)\n\n'
                                    '4.视频/图片目录: 可以将单个视频或单张图片放入该文件夹下，也可以是多张序列图放入该文件夹下\n\n'
                                    '5.背景视/图目录:可以将单个视频或单张图片放入该文件夹下，也可以是多张序列图放入该文件夹下\n\n'
                                    '6.Alpha输出目录: 只输出Alpha序列图，这里选择输出的目录')



if __name__ == '__main__':
    app =QApplication(sys.argv)
    bm=BackgroundMatting()
    bm.show()
    sys.exit(app.exec_())